/*
 * GMMReducer.h
 *
 *  Created on: 2011-8-18
 *      Author: yong
 */

#ifndef GMMREDUCER_H_
#define GMMREDUCER_H_
#include "Typedef.h"
#include "Reducer.h"
#include <iomanip>

namespace distrim
{
template<typename K, typename V>
class GMMReducer: public Reducer<K, V>
{
private:
	typedef boost::numeric::ublas::matrix<V> Matrix;
	typedef std::vector<Matrix> MV;
public:
	GMMReducer(const Config &config) :
		Reducer<K, V> (config), m_iterationCounter(0), m_maxIterCounter(
				config.GetMaxIterCount()),
				m_clusterNum(config.GetClusterNum()), m_totalPoints(
						config.GetTotalPoints()), m_valueElemCount(
						config.GetValueElemCount()), m_oldLikelihood(-1
						* numeric_limits<V>::max()), m_convergeDelta(
						config.GetConvergeDelta()), m_vecPartialSum(3
						* m_clusterNum + 1), m_vecParam(3 * m_clusterNum + 1),
				m_A(m_valueElemCount, m_valueElemCount)/*, m_B(m_valueElemCount,
	 m_valueElemCount)*/
	{
		BOOST_ASSERT(m_world.rank() == 0);
		setup();
	}

private:

	// Overloading.
	inline void setup()
	{
		// request space for parameters and intermediate partial sums.
		m_vecParam[0] = m_vecPartialSum[0] = Matrix(1, 1);
		for (size_t c = 1; c < m_clusterNum + 1; ++c)
		{
			m_vecParam[c] = m_vecPartialSum[c] = Matrix(1, 1);
			m_vecParam[c + m_clusterNum] = m_vecPartialSum[c + m_clusterNum]
					= Matrix(1, m_valueElemCount);
			m_vecParam[c + 2 * m_clusterNum] = m_vecPartialSum[c + 2
					* m_clusterNum]
					= Matrix(m_valueElemCount, m_valueElemCount);
		}
	}

	// Overloading.
	virtual void reduce()
	{
		// Reset partial sums.
		reset();
		m_world.barrier();
		// Reduce to get all partial sums.
		mpi::reduce(m_world, m_vecPartialSum, m_vecPartialSum, std::plus<MV>(),
				0);
	}

	// Overloading.
	inline void reset()
	{
foreach	(Matrix &m, m_vecPartialSum)
	{
		m.clear();
	}
}

// Overloading.
virtual bool terminate()
{
	if ((m_vecPartialSum[0](0, 0) - m_oldLikelihood) / m_totalPoints
			<= m_convergeDelta || m_iterationCounter >= m_maxIterCounter)
	{
		return true;
	}
	else
	{
		// Save for the next iteration to tell convergence.
		m_oldLikelihood = m_vecPartialSum[0](0, 0);
		return false;
	}
}

// Overloading.
virtual void update()
{
	m_vecParam[0] = m_vecPartialSum[0]; // New likelihood.
	std::cout << "Likelihood #" << m_iterationCounter << ": "
	<< std::setprecision(30) << m_vecParam[0](0, 0) << std::endl;
	for (size_t c = 1; c < m_clusterNum + 1; ++c)
	{
		m_vecParam[c] = m_vecPartialSum[c] / m_totalPoints; // New pi.
		m_vecParam[c + m_clusterNum] = m_vecPartialSum[c + m_clusterNum]
		/ m_vecPartialSum[c](0, 0); // New mu.
		prod(trans(m_vecParam[c + m_clusterNum]), m_vecParam[c + m_clusterNum], m_A);
		m_vecParam[c + 2 * m_clusterNum] = m_vecPartialSum[c + 2
		* m_clusterNum] / m_vecPartialSum[c](0,
				0) - m_A; // New sigma.
		/*
		 prod(trans(m_vecPartialSum[c + m_clusterNum]), m_vecParam[c
		 + m_clusterNum], m_A);
		 prod(trans(m_vecParam[c + m_clusterNum]), m_vecParam[c
		 + m_clusterNum], m_B);
		 m_vecParam[c + 2 * m_clusterNum] = (m_vecPartialSum[c + 2
		 * m_clusterNum] - m_A - trans(m_A)) / m_vecPartialSum[c](0,
		 0) + m_B; // New sigma.*/
	}
	m_world.barrier();
	for (int i = 1; i < m_world.size(); ++i)
	{
		// Send new parameters to mapper processes.
		m_world.send(i, msg_param_tag, m_vecParam);
	}
	++m_iterationCounter;
}
private:
mpi::communicator m_world;
size_t m_iterationCounter;
size_t m_maxIterCounter;
size_t m_clusterNum;
size_t m_totalPoints;
size_t m_valueElemCount;
V m_oldLikelihood;
V m_convergeDelta;
MV m_vecPartialSum;
MV m_vecParam;
Matrix m_A;
/*Matrix m_B;*/
};
}

#endif /* GMMREDUCER_H_ */
